#include "Dashun.h"
#include "Dashun_linux.hpp"


namespace qing {

    ///////////////////////////////////////////////////////////////////
    ///////////////////////////////////////////////////////////////////
    /*
    * 适用于机器人的DRL（深度强化学习）
    *
    * 蜘蛛机器人拥有10个舵机（4只胳膊、4只腿、2个摄像云台），它还有2个陀螺仪参数。
    *  将舵机目前的位置+陀螺仪参数传递给神经网络，输出10个舵机的概率值
    *  将概率值加上一个0~1之间的随机因子，获取最终输出
    *
    * 每个舵机180度，即 10 * 180 = 1800 个参数
    */
    class Test_DRL_20251207_1 {
    public:
        /* 构造函数，传入必要的数据 */
        Test_DRL_20251207_1 (int servo_num = 10):
         servo_num(servo_num){
            servo_all_angle = servo_num * ANGLE_MAX;    /* 计算所有舵机角度 */
            x_all_data = servo_all_angle + angle_num;   /* 计算所有的数据量 */
            MPU6050_init(); /* NOTE: 这两个模块应该遵从单例模式的使用方式 */
            PCA9685_init();
            fresh();
        }
        /* 刷新这次测试的数据 */
        void fresh() {
            init_nn();  /*  */
            servo_buffer = std::vector<int> (servo_num);
            angle_buffer = std::vector<float> (angle_num);
            x = std::vector<float> (x_all_data);
            r = x;
        }
        /* */
        void run_test(int t) {
            for (int i=0; i<t; ++i) {
                struct MPU6050_SensorData data = MPU6050_readSensorData();
                MPU6050_calAngle(data, angle_buffer[0], angle_buffer[1]);

                /* 奖励反向传播 */
                reward_func();

                /* 更新传感数据 */
                for (int j=0; j<angle_num; ++j) {
                    x[servo_all_angle+j] = angle_buffer[j];
                }

                /* 前向反馈 */
                for (int j=0; j<nn.size(); ++j) {
                    r = nn[j].forward(r);
                }

                /* 随机因子 */
                for (int j=0; j<servo_all_angle; ++j) {
                    r[j] += gen_rand();
                }

                /* 清空缓存 */
                for (int j=0; j<servo_num; ++j) {
                    servo_buffer[j] = 0;
                }

                /* 总结动作 */
                submit_actions();

                /* 执行... */
                actions_to_servo();
            }
        }

    private:
        std::vector<NeuralNetwork> nn;          /* 神经网络模型 */
        constexpr static int ANGLE_MAX = 180;   /* 舵机最大角度 */
        int servo_num;                          /* PCA9685舵机数量 */
        int angle_num = 2;                      /* MPU6050角度数量 */
        int servo_all_angle;                    /* 所有舵机角度 */
        int x_all_data;                         /* 所有的数据量 */
        std::vector<int>   servo_buffer;        /* 舵机角度缓冲区 */
        std::vector<float> angle_buffer;        /* 传感器角度 */
        std::vector<float> x;                   /* 输入数据 */
        std::vector<float> r;                   /* 输出数据 */

        
        /* 随机生成0~1 */
        static float gen_rand() {
            static std::random_device rd;
            static std::mt19937_64 engine(rd());
            static std::uniform_real_distribution<> dist(0.0, 1.0);
            return dist(engine);
        }

        /* 这个函数用于向nn中添加向量 */
        void init_nn() {
            nn = std::vector<NeuralNetwork> ();
            nn.push_back(NeuralNetwork::Create_in_Factory(x_all_data,       x_all_data * 2,     0.01, NeuralNetwork::ActivationFunc::Leaky_ReLU));
            nn.push_back(NeuralNetwork::Create_in_Factory(x_all_data * 2,   x_all_data * 2,     0.01, NeuralNetwork::ActivationFunc::Leaky_ReLU));
            nn.push_back(NeuralNetwork::Create_in_Factory(x_all_data * 2,   servo_all_angle,    0.01, NeuralNetwork::ActivationFunc::Sigmoid));
        }

        /*  */
        void reward_func() {
            ;
        }

        /* 总结动作 */
        void submit_actions() {
            for (int j =0; j<servo_num; ++j) {
                int maxidx  = -1;
                float max_v = -1; /* 现在x的最小值是-1 */
                for (int k=0; k<ANGLE_MAX; ++k) {
                    int idx = j*ANGLE_MAX+k;
                    int tt  = r[idx];
                    if (tt > max_v) {
                        max_v = tt;
                        maxidx = idx;
                    }
                }
                /* 结果 */
                servo_buffer[j] = maxidx;
            }
        }

        /* 执行动作 */
        void actions_to_servo() {
            for (int j=0; j<servo_num; ++j) {
                /* NOTE: 该函数将角度映射为0~180度 */
                PCA9685_setAngle(j, servo_buffer[j] / (float)ANGLE_MAX * 180.0);
            }
        }
    };
}